from sqlparse.tokens import Wildcard,DML,Keyword,Punctuation,Whitespace
from sqlparse.sql import Identifier, IdentifierList
from project_sqllineage.AstTree import AstTree

class ParseSQL:

    star_seq = 1  #星号出现的顺序，起始为1
    ast_table_list = []
    star_type_tag = False #判断* or t.* , False为* ,True为t.*
    select_index_max = -1 #全局变量 存放select对象在列表中最大位置的值
    subquery_alias = 'unkown' #20231018 子查询的别名
    seq_num = 1

    def parse_level(self,parsed,level,select_index,select_tag,from_tag,select_level):
        for item in parsed:
            if item.value == ' ':
                continue
            if item.ttype in DML and item.value.upper().strip() == 'SELECT':
                #select开始标志
                self.star_type_tag = False # 遇到select要初始化标志为False  (select t1.*,* 的情况 会有问题)
                select_tag = True
                from_tag = False
                select_level = level
                self.index_num = -1

                if select_index == self.select_index_max:
                    select_index = select_index + 1
                    self.select_index_max = select_index
                else:
                    select_index = self.select_index_max + 1
                    self.select_index_max = select_index

                ast_table = AstTree()
                ast_table.level = level
                self.ast_table_list.append(ast_table)
                self.ast_table_list[select_index].subquery_alias = self.subquery_alias #20231018
                self.ast_table_list[select_index].seq = self.seq_num
                self.seq_num += 1

            if item.ttype in Keyword and item.value.upper().strip() in ['FROM','JOIN','LEFT JOIN','RIGHT JOIN','FULL JOIN','INNER JOIN']:
                from_tag = True
                select_tag = False
                self.subquery_alias = 'unkown' #20231018

            if item.value.upper() in ['WHERE','GROUP BY','ORDER BY','UNION','UNION ALL']:
                # 结束标志
                from_tag = False
                select_tag = False
                self.subquery_alias = 'unkown' #20231018

            if item.is_group and 'SELECT' not in item.value.upper():
                for i in item:
                    if i.value.upper().strip() in ['WHERE','GROUP BY','ORDER BY','UNION','UNION ALL']:
                        from_tag = False
                        select_tag = False
                        self.subquery_alias = 'unkown'

            if item.value.upper().strip() in ['WHERE','GROUP BY','ORDER BY','UNION','UNION ALL']:
                from_tag = False
                select_tag = False
                self.subquery_alias = 'unkown'

            if item.value.upper().strip().split(' ')[0] == 'WHERE':
                from_tag = False
                select_tag = False
                self.subquery_alias = 'unkown'

            if isinstance(parsed,Identifier):
                self.star_type_tag = True

            if select_tag == True and from_tag == False:
                if select_level == level and (isinstance(item,IdentifierList) or isinstance(item,Identifier)):
                    if self.ast_table_list[select_index].column == 'unkown':
                        self.ast_table_list[select_index].column = item.value
                    else:
                        self.ast_table_list[select_index].column = self.ast_table_list[select_index].column + ' ' + item.value


                if not isinstance(item,IdentifierList) and level == select_level and item.ttype not in Keyword:
                    self.index_num += 1
                elif (level == select_level + 1) and item.ttype not in Punctuation and isinstance(item.parent,IdentifierList) and item.ttype != Whitespace:
                    self.index_num += 1

                #处理字段逻辑
                if item.ttype == Wildcard and self.star_type_tag == False:
                    #处理 select *
                    if parsed.value != '(*)':

                        #查看出现*号上一层是否是(*)
                        self.ast_table_list[select_index].index_num = self.index_num
                        self.ast_table_list[select_index].have_star = 'yes'
                        self.ast_table_list[select_index].column = item.value
                        self.ast_table_list[select_index].star_seq = self.star_seq
                        self.star_seq += 1

                elif item.ttype == Wildcard and self.star_type_tag == True:
                    #处理 select t.*
                    if parsed.value != '(*)':
                        self.ast_table_list[select_index].index_num = self.index_num
                        self.ast_table_list[select_index].have_star = 'yes'
                        self.ast_table_list[select_index].star_seq = self.star_seq
                        self.star_seq += 1

                elif isinstance(item, Identifier) or isinstance(item, IdentifierList) or item.is_group:
                    #其他情况需要递归下一层继续寻找*
                    self.parse_level(item,level+1,select_index,select_tag,from_tag,select_level)

            elif select_tag == False and from_tag == True:
                if item.ttype in Keyword and item.value.upper() == 'ON':
                    #出现on后面多为id=id，忽略此部分信息
                    from_tag = False
                    continue

                if isinstance(item,Identifier) or item.is_group:
                    if '(' in item.value:
                        table_name = 'subquery'
                        # print('table_name:'+table_name)
                        if hasattr(item, 'get_alias'):
                            if item.get_alias() != None:
                                alias_name = item.get_alias()
                                self.subquery_alias = alias_name #20231018
                            else:
                                alias_name = 'un_alias'
                        else:
                            alias_name = 'un_alias'

                        if level == select_level:
                            self.ast_table_list[select_index].table.append(table_name+'%%'+alias_name)

                        self.parse_level(item, level + 1 , select_index,select_tag,from_tag,select_level)
                    else:
                        if len(item.tokens) > 1 and item.tokens[1].value == '.':
                            table_name = item.tokens[0].value + item.tokens[1].value + item.tokens[2].value
                        else:
                            table_name = item.tokens[0].value ##需要抽象出一个取表名的函数
                        if hasattr(item, 'get_alias'):
                            if item.get_alias() != None:
                                alias_name = item.get_alias()
                            else:
                                alias_name = 'un_alias'
                        else:
                            alias_name = 'un_alias'

                        if level == select_level:
                            self.ast_table_list[select_index].table.append(table_name + '%%' + alias_name)

        return self.ast_table_list